import glob
import random
import os
import sys
import numpy as np
import torch
import torch.nn.functional as F
import torchvision.transforms as transforms
import json

from PIL import Image
from torch.utils.data import Dataset
from utils.augmentations import horisontal_flip


def parse_boxes(label, annotation_path: str):
    annotation_path = annotation_path.rstrip()
    boxes = []
    with open(annotation_path, 'r', encoding='gbk') as f:
        annotation = json.loads(f.read())['shapes']
        for item in annotation:
            boxes.append([label] + np.array(item['points']).reshape((4,)).tolist())
    return torch.tensor(boxes)


def pad_to_square(img, pad_value):
    c, h, w = img.shape
    dim_diff = np.abs(h - w)
    # (upper / left) padding and (lower / right) padding
    pad1, pad2 = dim_diff // 2, dim_diff - dim_diff // 2
    # Determine padding
    pad = (0, 0, pad1, pad2) if h <= w else (pad1, pad2, 0, 0)
    # Add padding
    img = F.pad(img, pad, "constant", value=pad_value)

    return img, pad


def resize(image, size):
    image = F.interpolate(image.unsqueeze(0), size=size, mode="nearest").squeeze(0)
    return image


def random_resize(images, min_size=288, max_size=448):
    new_size = random.sample(list(range(min_size, max_size + 1, 32)), 1)[0]
    images = F.interpolate(images, size=new_size, mode="nearest")
    return images


class ImageFolder(Dataset):
    def __init__(self, folder_path, img_size=416):
        self.files = sorted(glob.glob("%s/*.*" % folder_path))
        self.img_size = img_size

    def __getitem__(self, index):
        img_path = self.files[index % len(self.files)]
        # Extract image as PyTorch tensor
        img = transforms.ToTensor()(Image.open(img_path))
        # Pad to square resolution
        img, _ = pad_to_square(img, 0)
        # Resize
        img = resize(img, self.img_size)

        return img_path, img

    def __len__(self):
        return len(self.files)


class ListDataset(Dataset):
    def __init__(self, list_path, img_size=416, training=True, augment=True, multiscale=True, normalized_labels=False):
        with open(list_path, "r") as file:
            self.data_files = file.readlines()
            random.seed(0)
            random.shuffle(self.data_files)
        
        num = int(len(self.data_files) * 0.8)
        if training:
            self.data_files = self.data_files[:num]
        else:
            self.data_files = self.data_files[num:]

        with open('data/custom/support.json', 'r') as file:
            self.support_list = json.loads(file.read())

        with open('config/label_dict.json', 'r') as target:
            self.label_dict = json.loads(target.read())

        self.support_image = None
        self.img_size = img_size
        self.max_objects = 100
        self.augment = augment
        self.multiscale = multiscale
        self.normalized_labels = normalized_labels
        self.min_size = self.img_size - 3 * 32
        self.max_size = self.img_size + 3 * 32
        self.batch_count = 0

        self.create_support_image()

    def __getitem__(self, index):

        pack = self.data_files[index % len(self.data_files)].rstrip().split(',')
        label_name, annotation_path, img_path = pack[0], pack[-2], pack[-1]

        # ---------
        #  Image
        # ---------

        # Extract image as PyTorch tensor
        img = transforms.ToTensor()(Image.open(img_path).convert('RGB'))

        # Handle images with less than three channels
        if len(img.shape) != 3:
            img = img.unsqueeze(0)
            img = img.expand((3, img.shape[1:]))

        # Pad to square resolution
        img, pad = pad_to_square(img, 0)
        _, padded_h, padded_w = img.shape

        # ---------
        #  Label
        # ---------

        targets = None
        boxes = parse_boxes(self.label_dict[label_name], annotation_path)
        # Extract coordinates for image
        x1 = boxes[:, 1] + pad[0]
        y1 = boxes[:, 2] + pad[2]
        x2 = boxes[:, 3] + pad[0]
        y2 = boxes[:, 4] + pad[2]
        # Returns (x, y, w, h)
        boxes[:, 1] = ((x1 + x2) / 2) / padded_w
        boxes[:, 2] = ((y1 + y2) / 2) / padded_h
        boxes[:, 3] *= (x2 - x1) / padded_w
        boxes[:, 4] *= (y2 - y1) / padded_h

        targets = torch.zeros((len(boxes), 6))
        targets[:, 1:] = boxes

        # Apply augmentations
        if self.augment:
            if np.random.random() < 0.5:
                img, targets = horisontal_flip(img, targets)

        return img_path, img, targets
    
    def create_support_image(self):
        support_list = [None] * len(self.label_dict)
        for label_name, annotation_path, img_path in self.support_list:
            img = transforms.ToTensor()(Image.open(img_path).convert('RGB'))

            img, pad = pad_to_square(img, 0)

            box = parse_boxes(self.label_dict[label_name], annotation_path)[0]
            label = int(box[0].item())
            x1 = int(box[1]) + pad[0]
            y1 = int(box[2]) + pad[2]
            x2 = int(box[3]) + pad[0]
            y2 = int(box[4]) + pad[2]
            mask = torch.zeros((1,) + img.shape[-2:])
            mask[:, y1:y2+1, x1:x2+1] = 1.0
            img = torch.cat([img, mask], dim=0)
            support_list[label] = img
        self.support_image = torch.stack([resize(img, self.img_size) for img in support_list])

    def collate_fn(self, batch):
        paths, imgs, targets = list(zip(*batch))
        # Remove empty placeholder targets
        targets = [boxes for boxes in targets if boxes is not None]
        # Add sample index to targets
        for i, boxes in enumerate(targets):
            boxes[:, 0] = i
        targets = torch.cat(targets, 0)
        # Selects new image size every tenth batch
        if self.multiscale and self.batch_count % 10 == 0:
            self.img_size = random.choice(range(self.min_size, self.max_size + 1, 32))
        # Resize images to input shape
        imgs = torch.stack([resize(img, self.img_size) for img in imgs])
        self.batch_count += 1
        return paths, imgs, targets

    def __len__(self):
        return len(self.data_files)
